import sys
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QVBoxLayout, QPushButton, QFileDialog
from PyQt5.QtCore import Qt, QRect, QPoint, QSize, pyqtSignal
from PyQt5.QtGui import QPainter, QColor, QPen, QScreen, QPixmap, QImage
from PyQt5.Qt import QShortcut, QKeySequence


class ScreenshotWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.initShortcuts()

    def initUI(self):
        """初始化主窗口"""
        self.setWindowTitle("高级区域截图工具")
        self.setWindowFlags(Qt.WindowStaysOnTopHint)

        # 截图显示区域
        self.image_label = QLabel()
        self.image_label.setAlignment(Qt.AlignCenter)

        # 操作按钮
        self.new_btn = QPushButton("新建截图 (N)", self)
        self.save_btn = QPushButton("保存截图 (S)", self)
        self.copy_btn = QPushButton("复制到剪贴板 (C)", self)
        self.close_btn = QPushButton("关闭 (ESC)", self)

        # 禁用初始不可用的按钮
        self.save_btn.setEnabled(False)
        self.copy_btn.setEnabled(False)

        # 布局
        layout = QVBoxLayout()
        layout.addWidget(self.image_label)
        layout.addWidget(self.new_btn)
        layout.addWidget(self.save_btn)
        layout.addWidget(self.copy_btn)
        layout.addWidget(self.close_btn)
        self.setLayout(layout)

        # 连接信号
        self.new_btn.clicked.connect(self.startSelection)
        self.save_btn.clicked.connect(self.saveScreenshot)
        self.copy_btn.clicked.connect(self.copyToClipboard)
        self.close_btn.clicked.connect(self.close)

        # 截图相关变量
        self.selection_window = None
        self.screenshot = None

        self.resize(600, 400)

    def initShortcuts(self):
        """初始化快捷键"""
        QShortcut(QKeySequence("N"), self, self.startSelection)
        QShortcut(QKeySequence("S"), self, self.saveScreenshot)
        QShortcut(QKeySequence("C"), self, self.copyToClipboard)
        QShortcut(QKeySequence(Qt.Key_Escape), self, self.close)

    def startSelection(self):
        """开始选择截图区域"""
        self.hide()  # 隐藏主窗口

        # 创建全屏半透明选择窗口
        self.selection_window = SelectionWindow()
        self.selection_window.selectionCompleted.connect(self.onSelectionComplete)
        self.selection_window.showFullScreen()

    def onSelectionComplete(self, rect):
        """选择区域完成后的回调"""
        if rect.width() < 10 or rect.height() < 10:  # 忽略太小的选择
            self.show()
            return

        # 获取屏幕截图
        screen = QApplication.primaryScreen()
        self.screenshot = screen.grabWindow(0, rect.x(), rect.y(), rect.width(), rect.height())

        # 显示截图预览
        self.image_label.setPixmap(self.screenshot.scaled(
            self.image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

        # 启用相关按钮
        self.save_btn.setEnabled(True)
        self.copy_btn.setEnabled(True)

        self.show()

    def saveScreenshot(self):
        """保存截图到文件"""
        if not self.screenshot:
            return

        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存截图", "", "PNG 文件 (*.png);;JPEG 文件 (*.jpg *.jpeg)")

        if file_path:
            if file_path.endswith(('.png', '.PNG')):
                self.screenshot.save(file_path, 'PNG')
            else:
                self.screenshot.save(file_path, 'JPG', quality=90)

    def copyToClipboard(self):
        """复制截图到剪贴板"""
        if self.screenshot:
            QApplication.clipboard().setPixmap(self.screenshot)
            self.statusBar().showMessage("截图已复制到剪贴板", 2000)

    def closeEvent(self, event):
        """关闭事件处理"""
        if self.selection_window and self.selection_window.isVisible():
            self.selection_window.close()
        event.accept()


class SelectionWindow(QWidget):
    """区域选择窗口"""
    selectionCompleted = pyqtSignal(QRect)

    def __init__(self):
        super().__init__()
        self.initWindow()
        self.initVariables()

    def initWindow(self):
        """初始化窗口设置"""
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setWindowState(Qt.WindowFullScreen)
        self.setCursor(Qt.CrossCursor)

    def initVariables(self):
        """初始化选择变量"""
        self.start_pos = None
        self.end_pos = None
        self.is_selecting = False
        self.magnifier_enabled = True

        # 默认颜色设置
        self.selection_color = QColor(70, 150, 250, 150)
        self.border_color = Qt.white
        self.mask_color = QColor(0, 0, 0, 120)

    def paintEvent(self, event):
        """绘制选择区域和半透明遮罩"""
        painter = QPainter(self)

        # 绘制半透明背景
        painter.fillRect(self.rect(), self.mask_color)

        # 如果有选择区域，绘制选择框
        if self.start_pos and self.end_pos:
            selection_rect = QRect(self.start_pos, self.end_pos).normalized()

            # 绘制选择区域（完全透明）
            painter.setCompositionMode(QPainter.CompositionMode_Clear)
            painter.fillRect(selection_rect, Qt.transparent)
            painter.setCompositionMode(QPainter.CompositionMode_SourceOver)

            # 绘制选择区域边框
            painter.setPen(QPen(self.border_color, 2, Qt.DashLine))
            painter.drawRect(selection_rect)

            # 在角落显示尺寸信息
            painter.setPen(Qt.white)
            info_text = f"{selection_rect.width()} x {selection_rect.height()}"
            painter.drawText(selection_rect.bottomRight() + QPoint(-100, 20), info_text)

            # 绘制放大镜（如果启用）
            if self.magnifier_enabled and selection_rect.width() > 20 and selection_rect.height() > 20:
                self.drawMagnifier(painter, selection_rect)

    def drawMagnifier(self, painter, selection_rect):
        """绘制放大镜"""
        screen = QApplication.primaryScreen()
        screen_rect = screen.geometry()

        # 获取放大区域的屏幕坐标
        center_x = self.end_pos.x()
        center_y = self.end_pos.y()
        size = 80  # 放大镜大小

        # 调整坐标防止超出屏幕
        if center_x + size > screen_rect.right():
            center_x = screen_rect.right() - size
        if center_y + size > screen_rect.bottom():
            center_y = screen_rect.bottom() - size
        if center_x - size < screen_rect.left():
            center_x = screen_rect.left() + size
        if center_y - size < screen_rect.top():
            center_y = screen_rect.top() + size

        # 获取放大区域像素
        magnifier_rect = QRect(center_x - size // 2, center_y - size // 2, size, size)
        pixmap = screen.grabWindow(0, magnifier_rect.x(), magnifier_rect.y(), size, size)

        # 绘制放大镜背景
        painter.setBrush(QColor(0, 0, 0, 180))
        painter.setPen(Qt.NoPen)
        painter.drawEllipse(self.end_pos + QPoint(10, 10), 60, 60)

        # 绘制放大的像素
        painter.drawPixmap(
            self.end_pos.x() + 10 - 30,
            self.end_pos.y() + 10 - 30,
            pixmap.scaled(120, 120, Qt.KeepAspectRatio, Qt.FastTransformation))

        # 绘制中心十字线
        painter.setPen(QPen(Qt.red, 1))
        painter.drawLine(
            self.end_pos.x() + 10 - 10, self.end_pos.y() + 10,
            self.end_pos.x() + 10 + 10, self.end_pos.y() + 10)
        painter.drawLine(
            self.end_pos.x() + 10, self.end_pos.y() + 10 - 10,
            self.end_pos.x() + 10, self.end_pos.y() + 10 + 10)

    def mousePressEvent(self, event):
        """鼠标按下事件"""
        if event.button() == Qt.LeftButton:
            self.start_pos = event.pos()
            self.end_pos = event.pos()
            self.is_selecting = True
            self.update()

    def mouseMoveEvent(self, event):
        """鼠标移动事件"""
        if self.is_selecting:
            self.end_pos = event.pos()
            self.update()

    def mouseReleaseEvent(self, event):
        """鼠标释放事件"""
        if event.button() == Qt.LeftButton and self.is_selecting:
            self.end_pos = event.pos()
            self.is_selecting = False

            # 获取屏幕坐标系的选择区域
            screen_rect = QApplication.primaryScreen().geometry()
            global_start = self.mapToGlobal(self.start_pos)
            global_end = self.mapToGlobal(self.end_pos)
            selection_rect = QRect(global_start, global_end).normalized()

            # 确保选择区域在屏幕范围内
            selection_rect = selection_rect.intersected(screen_rect)

            self.selectionCompleted.emit(selection_rect)
            self.close()

    def keyPressEvent(self, event):
        """按键事件"""
        if event.key() == Qt.Key_Escape:
            self.close()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setStyle("Fusion")  # 设置统一的样式

    window = ScreenshotWindow()
    window.show()

    sys.exit(app.exec_())